"""
Library module for functions and classes related to traversing SSUMMO results
dictionaries, stored as a nested dictionary of dictionaries. These data types
are loadable directly from from the .pkl files, saved by SSuMMo.
"""

import os
from ssummo import CONFIG


def dict_walk(top, taxdict, topdown=True, random=True):
    """Generator function to recursively traverse a nested, hierarchical
    dictionary of taxonomic nodes. Yields tuples of (path name, node) for each
    node encounted in the tree.

    :param top: Path to prepend to each to each node.
    :param taxdict: Tree to traverse.
    :param topdown: If False, start by yielding leaf nodes and work backwards
                    up the tree.
    :param random: If False, sort the node names before yielding results.
    """
    if not random:
        names = sorted(taxdict.keys())
    else:
        names = taxdict.keys()
    if 'accessions' in names:
        del(names[names.index('accessions')])
    if topdown:
        yield (top, taxdict)
    for name in names:
        for x in dict_walk(os.path.join(top, name), taxdict[name],
                           topdown=topdown, random=random):
            yield x
    if not topdown:
        yield (top, taxdict)


def find_max_depth(node, depth=0, max_depth=0, deepest=[]):
    """Given a node, finds the deepest node and returns it's depth
    and  all nodes of thatdepth"""
    for branch in node.keys():
        if branch == 'accessions':
            continue
        if depth + 1 > max_depth:
            max_depth = depth + 1
            deepest = [branch]
        elif depth + 1 == max_depth:
            deepest.append(branch)
        (max_depth, deepest) = find_max_depth(node[branch], depth=depth+1,
                                        max_depth=max_depth, deepest=deepest)
    return (max_depth, deepest)


def find_start(tdict, options):
    """Locates the directory where to enter the SSUMMO loop.
    Default is to start in arbDBdir, which is configured in
    CONFIG.py.
    To change, give the command option '-start /some/path/to/dir'
    """
    if os.path.realpath(options['-start']) == os.path.realpath(CONFIG.arbDBdir):
        return tdict, CONFIG.arbDBdir
    else:
        startKeys = tdict.keys()
        if options['-start'].endswith( os.path.sep ):
            ## This is used in buildhmms to start with subdirectories of the
            ## given path, rather than starting there.
            appendPathSep = os.path.sep
        else:
            appendPathSep = ''
        startDir = options['-start'].rstrip( os.path.sep )
        pathList = startDir.split( os.path.sep)
        for key in startKeys:
            if key in pathList:
                firstNode = pathList.index(key)
                break
            else:
                continue
        node = tdict
        for nodeName in pathList[firstNode:]:
            if nodeName.strip() == '':
                continue
            else:
                node = node[nodeName]
                parentNode = nodeName
        startDir = os.path.join( CONFIG.arbDBdir, os.path.sep.join(pathList[firstNode:]) ) + appendPathSep
        options['-start'] = startDir 
    print("\n##\tStarting SSUMMO from node '{0}' at path '{1}'".format(parentNode, startDir))
    return (node, startDir)

def find_node_name(target_name, dict_node):
    """Searches the given tree (:data:`.dict_node`) for a node with the name
    :data:`.target_name`.

    :returns: the full path to the given node."""
    start_node = dict_node
    for path, node in dict_walk( '', start_node ):
        if target_name in node.keys():
            return path
    return None

def get_accessions(start_node, accessions=None):
    """Find and return all accessions from a given start_node."""
    for key in start_node.keys():
        if key == 'accessions':
#            for accession in start_node[key]:
#                if accession in accessions:
#                    print accession, start_node.keys()
            accessions += start_node['accessions'] 
        else:
            accessions = get_accessions(start_node[key], accessions=accessions)
    return accessions

def my_walk(top, taxdict, topdown=True, followlinks=False, random=False):
    """Traverses the dictionary taxdict and yields the full path
    appended to top, and a list of child nodes within the path
    yielded.
    """
    names = taxdict.keys()
    if not random:
        names = sorted( names )
    if 'accessions' in names:
        del( names[names.index('accessions')] )
    if  topdown:
        yield top, names 
    for name in names:
        path   =  os.path.join (top,  name )
        node = taxdict[name]
        for  x  in  my_walk(path, node,  topdown,  followlinks,random ):
            yield  x
    if  not  topdown:
        yield  top, names
